# -*- coding:utf-8 -*-
import os,sys
import re
import traceback
import time
import numpy as np
import talib
from datetime import datetime
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), os.pardir, os.pardir))
import supeanut_config
from CommonLib.mylog import mylog
from CommonLib.db.MongoDBTool import MongoDBTool


'''
作者：supeanut
创建时间：2016-xx-xx xx:xx:xx
功能描述：
	xxx
	xxxxx
相关配置：
	supeanut_config.XXX
历史改动：
	2016-xx-xx: xxxxxx
'''
class IndexCal:
	# init what
	def __init__(self):
		self.hq_series = ['open','close','high','low','vol_stock','vol_capital']
		self.date_series= ['date','date_year','date_month','date_day','date_quarter','date_week']
		self.index_series = ['kdj', 'ma', 'macd','rsi','obv','ema','cci','willr','atr','ad','adosc','adx',\
			'adxr','apo','aroon','aroonosc','bbands','bop','cdltocrows','cmo','dema','dx','ht_dcperiod',\
			'ht_dcphase','ht_phasor','ht_sine','ht_trendline','ht_trendmode','kama','mama','mfi','minus_dm',\
			'minus_di','mom','natr','plus_di','plus_dm','ppo','roc','rocp','rocr','sar','histhigh',\
			'volma_stock']
		self.index_param_dict = {\
		'date_year':{'cn_name':'年份','field':[],'default':[],'child':[]},\
		'date':{'cn_name':'日期','field':[],'default':[],'child':[]},\
		'date_month':{'cn_name':'月份','field':[],'default':[],'child':[]},\
		'date_day':{'cn_name':'日','field':[],'default':[],'child':[]},\
		'date_quarter':{'cn_name':'季度','field':[],'default':[],'child':[]},\
		'date_week':{'cn_name':'星期','field':[],'default':[],'child':[]},\
		'open':{'cn_name':'开盘价','field':[],'default':[],'child':[]},\
		'close':{'cn_name':'收盘价','field':[],'default':[],'child':[]},\
		'high':{'cn_name':'最高价','field':[],'default':[],'child':[]},\
		'low':{'cn_name':'最低价','field':[],'default':[],'child':[]},\
		'vol_stock':{'cn_name':'成交量（股）','field':[],'default':[],'child':[]},\
		'vol_capital':{'cn_name':'成交量（元￥）','field':[],'default':[],'child':[]},\
		'kdj':{'cn_name':'随机指标（stochastic）','field':['fastk_period', 'slowk_period', 'slowd_period'],'default':[9,3,3],'child':['k','d']},\
		'ma':{'cn_name':'均线','field':['timeperiod'],'default':[5],'child':[]},\
		'macd':{'cn_name':'指数平滑移动平均线','field':['fastperiod','slowperiod','signalperiod'],'default':[12,26,9],'child':['macd','macdsignal','macdhist']},\
		'rsi':{'cn_name':'相对强弱指标（Relative Strength Index）','field':['timeperiod'],'default':[14],'child':[]},\
		'obv':{'cn_name':'能量潮（On Balance Volume）','field':[],'default':[],'child':[]},\
		'ema':{'cn_name':'指数平均数指标（Exponential Moving Average）','field':['timeperiod'],'default':[30],'child':[]},\
		'cci':{'cn_name':'顺势指标（Commodity Channel Index）','field':['timeperiod'],'default':[14],'child':[]},\
		'willr':{'cn_name':'威廉指标（LarryWilliams Index）','field':['timeperiod'],'default':[14],'child':[]},\
		'atr':{'cn_name':'平均真实波幅（AverageTrueRange）','field':['timeperiod'],'default':[14],'child':[]},\
		'ad':{'cn_name':'ad','field':[],'default':[],'child':[]},\
		'adosc':{'cn_name':'adosc','field':['fastperiod','slowperiod'],'default':[3,10],'child':[]},\
		'adx':{'cn_name':'adx','field':['timeperiod'],'default':[14],'child':[]},\
		'adxr':{'cn_name':'adxr','field':['timeperiod'],'default':[14],'child':[]},\
		'apo':{'cn_name':'apo','field':['fastperiod','slowperiod'],'default':[12,26],'child':[]},\
		'aroon':{'cn_name':'aroon','field':['timeperiod'],'default':[14],'child':['aroondown','aroonup']},\
		'aroonosc':{'cn_name':'arronosc','field':['timeperiod'],'default':[14],'child':[]},\
		'bbands':{'cn_name':'bbands','field':['timeperiod','nbdevup','nbdevdn'],'default':[5,2,2],'child':['upperband','middleband','lowerband']},\
		'bop':{'cn_name':'bop','field':[],'default':[],'child':[]},\
		'cdltocrows':{'cn_name':'cdltocrows','field':[],'default':[],'child':[]},\
		'cmo':{'cn_name':'cmo','field':['timeperiod'],'default':[14],'child':[]},\
		'dema':{'cn_name':'dema','field':['timeperiod'],'default':[30],'child':[]},\
		'dx':{'cn_name':'dx','field':['timeperiod'],'default':[14],'child':[]},\
		'ht_dcperiod':{'cn_name':'ht_dcperiod','field':[],'default':[],'child':[]},\
		'ht_dcphase':{'cn_name':'ht_dcphase','field':[],'default':[],'child':[]},\
		'ht_phasor':{'cn_name':'ht_phasor','field':[],'default':[],'child':['inphase','quadrature']},\
		'ht_sine':{'cn_name':'ht_sine','field':[],'default':[],'child':['sine','leadsine']},\
		'ht_dcphase':{'cn_name':'ht_dcphase','field':[],'default':[],'child':[]},\
		'ht_trendline':{'cn_name':'ht_trendline','field':[],'default':[],'child':[]},\
		'ht_trendmode':{'cn_name':'ht_trendmode','field':[],'default':[],'child':[]},\
		'kama':{'cn_name':'kama','field':['timeperiod'],'default':[30],'child':[]},\
		'mama':{'cn_name':'mama','field':['fastlimit','fastlimit'],'default':[0.5,0.05],'child':['mama','fama']},\
		'mfi':{'cn_name':'mfi','field':['timeperiod'],'default':[14],'child':[]},\
		'minus_di':{'cn_name':'minus_di','field':['timeperiod'],'default':[14],'child':[]},\
		'minus_dm':{'cn_name':'minus_dm','field':['timeperiod'],'default':[14],'child':[]},\
		'mom':{'cn_name':'mom','field':['timeperiod'],'default':[10],'child':[]},\
		'natr':{'cn_name':'natr','field':['timeperiod'],'default':[14],'child':[]},\
		'plus_di':{'cn_name':'plus_di','field':['timeperiod'],'default':[14],'child':[]},\
		'plus_dm':{'cn_name':'plus_dm','field':['timeperiod'],'default':[14],'child':[]},\
		'ppo':{'cn_name':'ppo','field':['fastperiod','slowperiod'],'default':[12,26],'child':[]},\
		'roc':{'cn_name':'roc','field':['timeperiod'],'default':[10],'child':[]},\
		'rocr':{'cn_name':'rocr','field':['timeperiod'],'default':[10],'child':[]},\
		'rocp':{'cn_name':'rocp','field':['timeperiod'],'default':[10],'child':[]},\
		'sar':{'cn_name':'sar','field':['acceleration','maximum'],'default':[0.02,0.2],'child':[]},\
		'histhigh':{'cn_name':'histhigh','field':['histdays'],'default':[65535],'child':[]},\
		'volma_stock':{'cn_name':'volma_stock','field':['timeperiod'],'default':[5],'child':[]},\
		}
	
	def get_all_index_name(self, ):
		return self.index_series
	
	def get_all_hqdateindex_name(self, ):
		return self.date_series + self.hq_series + self.index_series
	
	def get_date_series_name(self, ):
		return self.date_series[1:]
		
	def get_param_field(self, index_name):
		return self.index_param_dict.get(index_name, False)
	
	def get_all_index_child(self,):
		child_list = []
		for index in self.index_param_dict:
			child_list += self.index_param_dict[index]['child']
		return child_list

	def dayoffset(self, ndarrays, dayoff):
		if dayoff == 0:
			return ndarrays
		pre_list = []
		for i in range(abs(dayoff)):
			pre_list.append(np.nan)
		pre_array = np.array(pre_list)
		if type(ndarrays) is tuple:
			a = []
			for ndarray in ndarrays:
				if dayoff < 0:
					ndarray_ = ndarray[:dayoff]
					a.append(np.append(pre_array, ndarray_))
				else:
					ndarray_ = ndarray[dayoff:]
					a.append(np.append(ndarray_, pre_array))
			return tuple(a)
		else:
			if dayoff < 0:
				ndarray = ndarrays[:dayoff]
				return np.append(pre_array, ndarray)
			else:
				ndarray = ndarrays[dayoff:]
				return np.append(ndarray,pre_array)


	# input_dict为list非nparray
	# 返回多list时，[[],[]]
	def cal_index(self, index_name, input_dict, params=[], dayoff=0, return_child = None):
		# 添加对date_year，close等支持
		# list --> ndarraay
		input_dict_np = {}
		date_time_list = input_dict['datetime']
		for key in input_dict:
			if key <> 'datetime':
				input_dict_np[key] = np.array(input_dict[key])
		if index_name == 'kdj':
			result = talib.STOCH(input_dict_np['high'], input_dict_np['low'], input_dict_np['close'], \
					fastk_period=params[0], slowk_period=params[1], slowd_period=params[2])
			if return_child == 'k':
				result = result[0]
			elif return_child == 'd':
				result = result[1]
		if index_name == 'macd':
			result = talib.MACD(input_dict_np['close'],fastperiod=params[0],slowperiod=params[1],signalperiod=params[2])
			if return_child == 'macd':
				result = result[0]
			elif return_child == 'macdsignal':
				result = result[1]
			elif return_child == 'macdhist':
				result = result[2]
		if index_name == 'mama':
			result = talib.MAMA(input_dict_np['close'],fastlimit=params[0],slowlimit=params[1])
			if return_child == 'mama':
				result = result[0]
			elif return_child == 'fama':
				result = result[1]
		if index_name == 'histhigh':
			histlen = params[0]
			input_dict['high']
		if index_name == 'minus_di':
			result = talib.MINUS_DI(input_dict_np['high'], input_dict_np['low'], input_dict_np['close'],timeperiod=params[0])
		if index_name == 'mom':
			result = talib.MOM(input_dict_np['close'],timeperiod=params[0])
		if index_name == 'sar':
			result = talib.SAR(input_dict_np['high'],input_dict_np['low'],acceleration=params[0],maximum=params[1])
		if index_name == 'plus_di':
			result = talib.PLUS_DI(input_dict_np['high'], input_dict_np['low'], input_dict_np['close'],timeperiod=params[0])
		if index_name == 'plus_dm':
			result = talib.PLUS_DM(input_dict_np['high'], input_dict_np['low'],timeperiod=params[0])
		if index_name == 'ppo':
			result = talib.PPO(input_dict_np['close'], fastperiod=params[0],slowperiod=params[1])
		if index_name == 'roc':
			result = talib.ROC(input_dict_np['close'],timeperiod=params[0])
		if index_name == 'rocr':
			result = talib.ROCR(input_dict_np['close'],timeperiod=params[0])
		if index_name == 'rocp':
			result = talib.ROCP(input_dict_np['close'],timeperiod=params[0])
		if index_name == 'natr':
			result = talib.NATR(input_dict_np['high'], input_dict_np['low'], input_dict_np['close'],timeperiod=params[0])
		if index_name == 'minus_dm':
			result = talib.MINUS_DM(input_dict_np['high'], input_dict_np['low'], timeperiod=params[0])
		if index_name == 'mfi':
			result = talib.MFI(input_dict_np['high'], input_dict_np['low'], input_dict_np['close'],\
				input_dict_np['vol_stock'],timeperiod=params[0])
		if index_name == 'ma':
			result = talib.MA(input_dict_np['close'], timeperiod=params[0])
		if index_name == 'volma_stock':
			result = talib.MA(input_dict_np['vol_stock'], timeperiod=params[0])
		if index_name == 'dema':
			result = talib.DEMA(input_dict_np['close'], timeperiod=params[0])
		if index_name == 'dx':
			result = talib.DX(input_dict_np['high'],input_dict_np['low'],input_dict_np['close'],timeperiod=params[0])
		if index_name == 'cmo':
			result = talib.CMO(input_dict_np['close'], timeperiod=params[0])
		if index_name == 'atr':
			result = talib.ATR(input_dict_np['high'],input_dict_np['low'],input_dict_np['close'],timeperiod=params[0])
		if index_name == 'willr':
			result = talib.WILLR(input_dict_np['high'],input_dict_np['low'],input_dict_np['close'],timeperiod=params[0])
		if index_name == 'cci':
			result = talib.CCI(input_dict_np['high'],input_dict_np['low'],input_dict_np['close'],timeperiod=params[0])
		if index_name == 'ema':
			result = talib.EMA(input_dict_np['close'],timeperiod=params[0])
		if index_name == 'obv':
			result = talib.OBV(input_dict_np['close'],input_dict_np['vol_stock'])
		if index_name == 'rsi':
			result = talib.RSI(input_dict_np['close'],timeperiod=params[0])
		if index_name == 'ht_dcperiod':
			result = talib.HT_DCPERIOD(input_dict_np['close'])
		if index_name == 'ht_dcphase':
			result = talib.HT_DCPHASE(input_dict_np['close'])
		if index_name == 'ht_phasor':
			result = talib.HT_PHASOR(input_dict_np['close'])
			if return_child == 'inphase':
				result = result[0]
			elif return_child == 'quadrature':
				result = result[1]
		if index_name == 'ht_sine':
			result = talib.HT_SINE(input_dict_np['close'])
			if return_child == 'sine':
				result = result[0]
			elif return_child == 'leadsine':
				result = result[1]
		if index_name == 'ht_trendline':
			result = talib.HT_TRENDLINE(input_dict_np['close'])
		if index_name == 'ht_trendmode':
			result = talib.HT_TRENDMODE(input_dict_np['close'])
		if index_name == 'kama':
			result = talib.KAMA(input_dict_np['close'],timeperiod=params[0])
		if index_name == 'close':
			result = input_dict_np['close']
		if index_name == 'open':
			result = input_dict_np['open']
		if index_name == 'high':
			result = input_dict_np['high']
		if index_name == 'low':
			result = input_dict_np['low']
		if index_name == 'vol_stock':
			result = input_dict_np['vol_stock']
		if index_name == 'vol_capital':
			result = input_dict_np['vol_capital']
		if index_name == 'ad':
			result = talib.AD(input_dict_np['high'],input_dict_np['low'],input_dict_np['close'],\
								input_dict_np['vol_stock'])
		if index_name == 'adosc':
			result = talib.ADOSC(input_dict_np['high'],input_dict_np['low'],input_dict_np['close'],\
				input_dict_np['vol_stock'],fastperiod=params[0],slowperiod=params[1])
		if index_name == 'adx':
			result = talib.ADX(input_dict_np['high'],input_dict_np['low'],input_dict_np['close'],\
				timeperiod=params[0])
		if index_name == 'adxr':
			result = talib.ADXR(input_dict_np['high'],input_dict_np['low'],input_dict_np['close'],\
				timeperiod=params[0])
		if index_name == 'apo':
			result = talib.APO(input_dict_np['close'],fastperiod=params[0],slowperiod=params[1])
		if index_name == 'aroon':
			result = talib.AROON(input_dict_np['high'],input_dict_np['low'],timeperiod=params[0])
			if return_child == 'aroondown':
				result = result[0]
			elif return_child == 'aroonup':
				result = result[1]
		if index_name == 'bbands':
			result = talib.BBANDS(input_dict_np['close'],timeperiod=params[0],nbdevup=params[1],nbdevdn=params[2])
			if return_child == 'upperband':
				result = result[0]
			elif return_child == 'middleband':
				result = result[1]
			elif return_child == 'lowerband':
				result = result[2]
		if index_name == 'bop':
			result = talib.BOP(input_dict_np['open'],input_dict_np['high'],input_dict_np['low'],input_dict_np['close'])
		if index_name == 'aroonosc':
			result = talib.AROONOSC(input_dict_np['high'],input_dict_np['low'],timeperiod=params[0])
		if index_name == 'cdltocrows':
			result = talib.CDL2CROWS(input_dict_np['open'],input_dict_np['high'],input_dict_np['low'],input_dict_np['close'])
		if index_name == 'date':
			result = [k.split(' ')[0] for k in date_time_list]
			return result
		if index_name == 'date_year':
			result = [(k.split(' ')[0]).split('-')[0] for k in date_time_list]
			return result
		if index_name == 'date_month':
			result = [(k.split(' ')[0]).split('-')[1] for k in date_time_list]
			return result
		if index_name == 'date_day':
			result = [(k.split(' ')[0]).split('-')[2] for k in date_time_list]
			return result
		if index_name == 'date_quarter':
			result = []
			for date in date_time_list:
				quarter = ''
				month = date.split(' ')[0].split('-')[1]
				if month in ['01','02','03']:
					quarter = 'quarter1'
				if month in ['04','05','06']:
					quarter = 'quarter2'
				if month in ['07','08','09']:
					quarter = 'quarter3'
				if month in ['10','11','12']:
					quarter = 'quarter4'
				result.append(quarter)
			return result
		if index_name == 'date_week':
			result = []
			for date in date_time_list:
				date_dt = datetime.strptime(date.split(' ')[0], '%Y-%m-%d')
				weekday_i = date_dt.weekday()
				weekday_str = ""
				if weekday_i == 0:
					weekday_str = "Monday"
				elif weekday_i == 1:
					weekday_str = "Tuesday"
				elif weekday_i == 2:
					weekday_str = "Wednesday"
				elif weekday_i == 3:
					weekday_str = "Thursday"
				elif weekday_i == 4:
					weekday_str = "Friday"
				elif weekday_i == 5:
					weekday_str = "Saturday"
				elif weekday_i == 6:
					weekday_str = "Sunday"
				result.append(weekday_str)
			return result
		# offset day
		return_result = self.dayoffset(result,dayoff)
		# ndarray-->list
		return_data = []
		if isinstance(return_result,tuple):
			for result in return_result:
				return_data.append(result.tolist())
		else:
			return_data.extend(return_result.tolist())
		return return_data

		#one_list = np.random.random(50)
		#print one_list
		# --MA timeperiod>=1,int
		#print talib.MA(one_list, timeperiod=3)
		# --matype 0=SMA, 1=EMA, 2=WMA, 3=DEMA, 4=TEMA, 5=TRIMA, 6=KAMA, 7=MAMA
		#result = talib.MA(one_list, timeperiod=10, matype=8)
		# --T3  0.0<vfactor<1.0,double
		#print talib.T3(one_list, timeperiod=10, vfactor=0.2)



if __name__ == '__main__':
	obj = IndexCal()
	inputs = {
		'open': np.random.random(100),
		'high': np.random.random(100),
		'low': np.random.random(100),
		'close': np.random.random(100),
		'vol_stock': np.random.random(100),
		'datetime': np.random.random(100),
						}
	'''
	result = obj.cal_index(index_name='kdj', input_dict=inputs, params=[5,2,2], dayoff=1, return_child='k')
	print result
	print type(result)
	result = obj.cal_index(index_name='macd', input_dict=inputs, params=[12,26,9], dayoff=0, return_child=None)
	print result
	'''
	result = obj.cal_index(index_name='cdltocrows', input_dict=inputs, params=[], dayoff=0,return_child=None)
	print result
